import gradio as gr
import numpy as np
import pickle
from PIL import Image

with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)

def predict_digit(image):
    # 检查image是否为字典，如果是，则提取图像数据
    if isinstance(image, dict) and "image" in image:
        image = image["image"]
    
    # 确保输入是 NumPy 数组
    image = np.array(image)
    
    # 将输入的手写图像转换为 PIL 格式的灰度图，并调整大小为28x28
    image = Image.fromarray(image).convert('L').resize((8, 8))
    
    # 将图像数据转换为数组，并将其展平为1D数组
    img_array = np.array(image).reshape(1, -1)
    
    img_array = (img_array /255 )*16# 将像素值缩放到0-1之间
    
    # 进行预测
    prediction = knn_model.predict(img_array)
    return int(prediction[0])

interface = gr.Interface(
    fn=predict_digit,  # 调用的函数
    inputs=gr.Sketchpad(),  # 手写板输入
    outputs=gr.Label(),  # 输出是预测的标签
    title="手写数字识别",  # Web 应用标题
    description="在手写板上输入一个数字，模型将返回预测结果。"  # 描述信息
)

interface.launch()